import pandas as pd
from StockTradingData import StockTradingData
from queue import Queue

from libs.Objects import Objects
from libs.StockBase import StockBase

DAY_FORMAT = '%Y-%m-%d'


class IndexCollector(object):
    def __init__(self):
        StockBase.reConnect()

    def queryShSzAmount(self, tradingDay):
        # sz.399106 深证综合指数
        # sz.399001 深证成份指数(价格)
        # sh.000001 上证A股指数
        queue = Queue()
        # mutilProcess
        for code in ['sh.000001', 'sz.399001']:
            StockTradingData(tradingDay, code, queue).download()

        print('sh.000001 sz.399001 finished! ' + str(queue.qsize()))
        df_arry = [queue.get() for i in range(queue.qsize())]
        all_df = pd.concat(df_arry)

        all_df['amount_len'] = pd.to_numeric(all_df['amount'], errors='coerce')
        return all_df[['amount_len']].sum().sum()


if __name__ == '__main__':
    START_DAY = '2025-04-09'
    END_DAY = '2025-04-10'
    for i in range(100):
        curday = Objects.getDiffDay(START_DAY, i)
        if curday == END_DAY:
            exit()
        if not StockBase().isTradingDate(curday):
            continue
        a = IndexCollector().queryShSzAmount(curday)
        print(a)
